#!/usr/bin/python

""" Generate datas for cassandra 0.7.4
argv[1]: keyspace name
argv[2]: number of Encode_Entity
"""

import os, sys, time
import pycassa
import pycassa.system_manager
from google.appengine.ext import db
from google.appengine.api import users
import numpy

from protoBuff_decode import *

class Greeting(db.Model):
  author = db.UserProperty()
  content = db.StringProperty(multiline=True)
  date = db.DateTimeProperty(auto_now_add=True)

if __name__ == '__main__':
    keyspace_name = sys.argv[1]
    data_num = int(sys.argv[2])
    
    cassandra_7_srv_addr = '128.111.55.223'
    cassandra_7_port = 29160
    column_family_name = "Standard1"
    
    print '======================================================================'
    print 'Connecting to ' + cassandra_7_srv_addr + ':' + str(cassandra_7_port)
    print 'Using Keyspace : ' + keyspace_name + ', Column Family : ' + column_family_name
    cass_7_pool = pycassa.connect(keyspace_name, [cassandra_7_srv_addr + ':' + str(cassandra_7_port)])
    sys_manage = pycassa.system_manager.SystemManager(cassandra_7_srv_addr + ':' + str(cassandra_7_port))
    print 'Connected!'
    
    # create keyspace_name if does not exist
    keyspace_list = sys_manage.list_keyspaces()
    if keyspace_name not in keyspace_list:
        print 'Keyspace "' + keyspace_name + '" does NOT exist, create it'
        sys_manage.create_keyspace(keyspace_name, 1)
    
    # create a new column_family_dict 
    column_family_dict = sys_manage.get_keyspace_column_families(keyspace_name)
    if column_family_dict.has_key(column_family_name):
        print 'Drop column family "' + column_family_name + '"'
        sys_manage.drop_column_family(keyspace_name, column_family_name)
    print 'Create column family "' + column_family_name + '"'
    sys_manage.create_column_family(keyspace_name, column_family_name)
    
    cass_7_col_fam = pycassa.ColumnFamily(cass_7_pool, column_family_name)    
    print '======================================================================'

    # for google appengine api to be used
    os.environ['AUTH_DOMAIN'] = 'ucsb.edu'
    os.environ['APPLICATION_ID'] = 'bear_app'
    
    print 'Generating ' + str(data_num) + ' Greetings ...'
    
    zipf_int_list = numpy.random.zipf(2.0, data_num) # zipf distribution
    
    for i in range(data_num):
        if i%10000 == 0:
            print 'Greeting[' + str(i) + '] : '
        greeting = Greeting()
        user_name = 'user' + str(zipf_int_list[i])
        greeting.author = users.User(user_name)
        greeting.content = 'I\'m ' + user_name + ' at ' + str(time.time())
        
        # encode by protocol buffer
        proto_buf = db.model_to_protobuf(greeting)
        bin_str = proto_buf.Encode()
        
        # put into cassandra
        Txn_Num = str(i * 2 + 1)
        
        key_str = 'JOURNAL_____guestbook___1/guestbook/Greeting:' + str((i + 1) * 10).rjust(64, '0') + '/' + Txn_Num.rjust(64, '0')
        cass_7_col_fam.insert(key_str, {'Encoded_Entity' : bin_str})
        
        key_str = 'guestbook___Greeting___1/guestbook/Greeting:' + str((i + 1) * 10).rjust(64, '0')
        cass_7_col_fam.insert(key_str, {'Encoded_Entity' : bin_str})
        cass_7_col_fam.insert(key_str, {'Txn_Num' : Txn_Num})
    
    print 'Done!'
    
